package com.element.oauth2.server.global.provider;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.security.authentication.AuthenticationProvider;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.AuthenticationException;
import org.springframework.security.oauth2.core.OAuth2AuthenticationException;
import org.springframework.security.oauth2.core.OAuth2Token;
import org.springframework.security.oauth2.server.authorization.OAuth2Authorization;
import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationService;
import org.springframework.security.oauth2.server.authorization.OAuth2TokenType;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2ClientAuthenticationToken;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2TokenRevocationAuthenticationToken;
import org.springframework.security.oauth2.server.authorization.client.RegisteredClient;
import org.springframework.util.Assert;

public class OAuthTokenRevocationAuthenticationProvider implements AuthenticationProvider {

    private final Log logger = LogFactory.getLog(this.getClass());
    private final OAuth2AuthorizationService authorizationService;

    public OAuthTokenRevocationAuthenticationProvider(OAuth2AuthorizationService authorizationService) {
        Assert.notNull(authorizationService, "authorizationService cannot be null");
        this.authorizationService = authorizationService;
    }

    public Authentication authenticate(Authentication authentication) throws AuthenticationException {
        OAuth2TokenRevocationAuthenticationToken tokenRevocationAuthentication = (OAuth2TokenRevocationAuthenticationToken) authentication;
        OAuth2ClientAuthenticationToken clientPrincipal = getAuthenticatedClientElseThrowInvalidClient(tokenRevocationAuthentication);
        RegisteredClient registeredClient = clientPrincipal.getRegisteredClient();
        OAuth2TokenType tokenType = null;
        if (OAuth2TokenType.ACCESS_TOKEN.getValue().equals(tokenRevocationAuthentication.getTokenTypeHint())) {
            tokenType = OAuth2TokenType.ACCESS_TOKEN;
        } else if (OAuth2TokenType.REFRESH_TOKEN.getValue().equals(tokenRevocationAuthentication.getTokenTypeHint())) {
            tokenType = OAuth2TokenType.REFRESH_TOKEN;
        }
        OAuth2Authorization authorization = this.authorizationService.findByToken(tokenRevocationAuthentication.getToken(), tokenType);
        if (authorization == null) {
            if (this.logger.isTraceEnabled()) {
                this.logger.trace("Did not authenticate token revocation request since token was not found");
            }

            return tokenRevocationAuthentication;
        } else if (null == registeredClient || !authorization.getRegisteredClientId().equals(registeredClient.getId())) {
            throw new OAuth2AuthenticationException("invalid_client");
        } else {
            OAuth2Authorization.Token<OAuth2Token> token = authorization.getToken(tokenRevocationAuthentication.getToken());
            this.authorizationService.remove(authorization);
            if (this.logger.isTraceEnabled()) {
                this.logger.trace("Saved authorization with revoked token");
                this.logger.trace("Authenticated token revocation request");
            }
            assert token != null;
            return new OAuth2TokenRevocationAuthenticationToken(token.getToken(), clientPrincipal);
        }
    }

    public boolean supports(Class<?> authentication) {
        return OAuth2TokenRevocationAuthenticationToken.class.isAssignableFrom(authentication);
    }

    private OAuth2ClientAuthenticationToken getAuthenticatedClientElseThrowInvalidClient(Authentication authentication) {
        OAuth2ClientAuthenticationToken clientPrincipal = null;
        if (OAuth2ClientAuthenticationToken.class.isAssignableFrom(authentication.getPrincipal().getClass())) {
            clientPrincipal = (OAuth2ClientAuthenticationToken) authentication.getPrincipal();
        }

        if (clientPrincipal != null && clientPrincipal.isAuthenticated()) {
            return clientPrincipal;
        } else {
            throw new OAuth2AuthenticationException("invalid_client");
        }
    }
}